[experimental] rocSHMEM as an allocation provider - #550

allocate_symmetric() cached the table built from the first allocation and
returned it for every later one. Translation still worked, because the per-peer
offset is constant across the heap, but peer_bases[local_rank] was the first
allocation's base rather than the current tensor's -- contradicting the stated
invariant and breaking the assertion in tests/manual_rocshmem_provider.py, which
checks it on a second allocation.

Cache the per-peer offsets instead and build each allocation's table from its
own base. peer_bases[local_rank] is now that allocation's base for every
allocation, while the shared offsets keep a table from one allocation able to
translate another's pointers, which iris.copy relies on.

This also removes a duplicated rocshmem_ptr sweep: the first allocation
previously queried every peer twice, once to build the map and once to seed the
cache. It is now queried once per process.

peer_bases is created on tensor.device rather than a device captured when the
provider was constructed, so the table cannot end up on a different device than
the memory it describes.

test_table_is_context_wide asserted the two tables were equal, which no longer
holds and was the weaker property anyway. It is now test_peer_offsets_are_shared
and checks what actually matters: each table's local entry is its own
allocation's base, and the per-peer offsets agree. Unreachable peers are
excluded from that comparison, since their entry is 0 rather than base + offset.

Verified on 2 ranks: pytest 4 passed, manual test PASS including the
cross-allocation check.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Review follow-ups, all documentation:

- allocate_symmetric: say it is collective, describe the returned table's
  shape, dtype, device and invariant, and point at allocate_symmetric_map
  for the direct mask.
- symmetric_address_map: note it builds a fresh table per call, why that is
  once per allocation rather than per launch, and why it is deliberately not
  memoised by data_ptr.
- free: explain it cannot be automated. rocshmem_free is documented as
  collective and must be called by all PEs, so a __del__ or weakref finalizer
  would let ranks diverge on GC timing and hang.
- module: record that rocshmem4py is a standalone package that statically
  links rocSHMEM rather than linking it at run time, so its USE_IPC setting
  is fixed at its build time.

The provider tests have only ever skipped in CI, because rocshmem4py was not
installed. Add it, so they actually execute.

Installing it is a source build but a short one. There is no prebuilt
rocshmem4py wheel on any index -- not PyPI, not the ROCm nightly indexes -- and
the bindings do not build rocSHMEM themselves (find_package(rocshmem 3.5.0
CONFIG REQUIRED), no FetchContent), so rocSHMEM is built first and the bindings
are pip-installed against it with CMAKE_PREFIX_PATH.

Scope is IPC only, which is what the provider uses and all a single-node runner
can exercise. Upstream already defaults USE_IPC=ON and USE_GDA=OFF, so no
conduit flags are passed, which keeps MPI and the RDMA provider libraries out of
it. Built for gfx942 to match the MI325X runners.

One installer shared by both image definitions rather than two copies. That
required widening the Docker build context from docker/ to the repo root so the
Dockerfile can COPY it, hence the .dockerignore; the Apptainer def pulls the same
file in via %files.

Also adds iris/experimental/README.md covering install, verification and the
collective-call contract, as requested in review.

Both are pip installs from source, but the README documented the git+ form while
the CI script used a local path, with nothing explaining the difference.

One checkout serves the core build and the bindings, so they are guaranteed to be
the same revision. A git+ URL would have pip clone the monorepo again at whatever
develop is at by then, and find_package would not catch the skew because it only
compares versions while the bindings statically link the core.

The docstring claimed the same signature and return shape as
Iris.allocate_symmetric. Iris has no such method -- its allocation API is
zeros/ones/full/uniform/as_symmetric, with a context-wide table from
get_heap_bases(). The (tensor, peer_bases) shape comes from the provider
interface being proposed, not from an existing method.

Comment only. The CI bases are ROCm 7.2.1 (apptainer) and 7.1 (docker), both
older than the 7.14 artifacts that ship rocSHMEM's static library and headers,
so the source build is still required. Records where to cut it when a base image
bumps, and that rocshmem4py stays a source build either way until its TheRock
packaging lands.

The provider tests now run in CI rather than skip, and rocSHMEM aborts inside
init with no message: a bare SIGABRT through library_init -> IPCBackend ->
Backend::init, 40ms in, with nothing to say why.

Set ROCSHMEM_DEBUG_LEVEL=info in both images so it prints its config banner, the
env vars it actually saw, and the backend it selected. Also write the rank, world
size and visible GPU count straight to fd 2 before the call, which survives both
pytest's capture and the abort, so the crash says what it was attempting.

Diagnostic only; drop both once the tests pass.

Review feedback.

The rationale for the source build, and for ROCSHMEM_DEBUG_LEVEL, was spelled out
in install_rocshmem.sh and then again in both image definitions. Keep it in the
script and leave a pointer at each call site.

Mark the ROCm 7.14 note as a TODO so it shows up when someone greps for work to
drop, rather than reading as background.

The previous diagnostic round produced nothing: neither rocSHMEM's banner nor a
write to fd 2 reached the CI log. pytest captures at the fd level and has already
redirected fd 2 by the time a fixture runs, so both landed in a buffer that is
discarded when the process aborts. faulthandler's output survived only because it
dups the original fd 2 at interpreter startup, which is why "Fatal Python error"
came through and nothing else did.

Suspend capture around init via the capturemanager plugin so rocSHMEM's own
logging, and the rank/GPU-count line, reach the log.

Also fold two overlapping "Returns (tensor, peer_bases)" paragraphs in
allocate_symmetric's docstring into one; they came from two separate edits.

Diagnostic; drop with the debug level once the tests pass.

With capture suspended, rocSHMEM's banner finally reached the log and named the
problem outright:

  # Compiled Arch(s)         : gfx942
  # System Arch              : gfx950:sramecc+:xnack-
  # System Arch is supported : No

The runner label is linux-mi325-8gpu-ossci-rad, and MI325X is gfx942, but the
hardware reports gfx950. rocSHMEM's device code has to match the GPU it runs on;
built for the wrong arch, the loaded code object has none of its device globals
and init aborts in HIP:

  hip_global.cpp:70 : Cannot create GlobalVar Obj for symbol:
  _ZN8rocshmem14logd_constantsE

Build both arches rather than swapping one guess for another, since the label
cannot be trusted to track the pool.

They did their job: the banner named the arch mismatch that 578394b fixes.

Removes ROCSHMEM_DEBUG_LEVEL=info from both images, and the fd-2 write plus the
pytest capture suspension from the provider fixture. The debug level printed a
config banner on every run, and the capture suspension is scaffolding that should
not outlive the bug it was added for.

Kept as its own commit so `git revert` restores the whole apparatus in one step
if a rocSHMEM init failure ever needs diagnosing again. Note that without it such
a failure is mute: pytest captures at the fd level, so rocSHMEM's output is
buffered and discarded when the process aborts.

ROCm#549 landed it on main, returning the same (tensor, peer_bases) pair with
peer_bases[cur_rank] == data_ptr(). An earlier commit dropped this reference
because the method did not exist yet; matching it is the whole point of the
provider, so say so.

$HOME is shared across the runners, so every job in a run uses one cache
directory. Jobs that start together all find no image and all build concurrently
into the same path -- observed in run 35023261876, where jobs on iris-mi350x-0
and iris-mi350x-2 both reported "Image or checksum not found" and built, while
jobs on -1, -2 and -3 read a checksum one of them had written.

Two problems with that. `apptainer build --force` writes IMAGE_PATH in place, so
a builder truncates the image other jobs are currently executing. And every job
that loses the race still pays for a full redundant build.

Take an flock around the check-and-build, and re-check inside it so whoever waits
usually finds the image already there. Build to a private temp path and rename
into place: rename is atomic and leaves the inode alone, so a job already running
the old image is unaffected.

Verified with two concurrent builders against a stub: one build, the other waited
and reused it, no temp files left. On build failure: exit 1, no checksum written,
no image published, temp removed.